##################################################
# (1) init ==> B <= 0
# (2) unsafe ==> B > 0
# (3) domain /\ x /\ x' ==> B(x') - (1-alpha) * B(x) < 0
##################################################
# max(B(x))  s.t. x in init
# if max(B)<=0: condition (1) is OK
import math
import time
import cvxpy as cp
import dccp
import gurobipy as gp
from gurobipy import GRB

import verify_dccp.getMathExpress2 as express
from prob import *
import numpy as np
import verify_piecewise_utils
from NN import cdinn2
import torch


def verif(model):
    cex_init = []
    res_init = []
    cex_unsafe = []
    res_unsafe = []
    cex_domain = []
    res_domain = []

    x = cp.Variable(DIM)  # x0, x1, x2
    v = cp.Variable()

    f1, f2 = express.getExpr(x, model)

    # print("curvature of f1:", f1.curvature)
    # print("curvature of f2:", f2.curvature)

    print("=====================X_INIT=========================")
    objective = cp.Maximize(f1 - v)
    constraints = [v == f2]

    constraints += [cp.square(x[0] - 1.5) + cp.square(x[1] - 0.5) + cp.square(x[2] - 1.5) <= 0.25]

    problem = cp.Problem(objective, constraints)
    # print("problem is DCCP:", dccp.is_dccp(problem))
    result = problem.solve(method='dccp')
    # print("x.value:", x.value)  # [ 1.34102711 -0.47405446]
    # print("v.value:", v.value)  # 4.158081062071981
    # print("status:", problem.status) # Converged
    # print("cost value0 =", result[0])  # -1.161325675958789
    if result[0] > 0.0:
        cex_init.append(x.value)
        res_init.append(result[0])
    else:
        print("init is safe")

    print("=====================X_UNSAFE=========================")
    objective = cp.Minimize(v - f2)
    constraints = [v == f1]

    constraints += [cp.square(x[0] - 0.5) + cp.square(x[1] - 1.5) + cp.square(x[2] - 1.5) <= 0.25]

    problem = cp.Problem(objective, constraints)
    # print("problem is DCCP:", dccp.is_dccp(problem))
    # x.value = np.array([1, 2])
    result = problem.solve(method='dccp')
    # print("x.value:", x.value)
    # print("v.value:", v.value)
    # print("status:", problem.status)
    # print("cost value0 =", result[0])
    if result[0] < 0.0:
        cex_unsafe.append(x.value)
        res_unsafe.append(result[0])
    else:
        print("unsafe is safe")

    print("=====================X_DOMAIN=========================")
    fw = open('log/verify_c14.txt', 'w+', encoding='utf-8')

    data_dir = "./data/"
    # datas = np.loadtxt(data_dir + "10points_0.1.txt", dtype=np.float64).reshape([-1, 8])
    datas = np.loadtxt(data_dir + "-2_2_0.5_20_3d.txt", dtype=np.float64).reshape([-1, 27])
    # datas = np.loadtxt(data_dir + "-2_2_0.1_10.txt", dtype=np.float64).reshape([-1, 4])
    # (x_l, x_r) + (y_l, y_r) + (w0, w1) + (b_l, b)

    var_x0_b = cp.Variable()
    var_x1_b = cp.Variable()
    var_x2_b = cp.Variable()
    var_x01_b = cp.Variable()
    var_x02_b = cp.Variable()
    var_x12_b = cp.Variable()

    x0_min = domain_min[0]  # -2
    x0_max = domain_max[0]  # 2
    x1_min = domain_min[1]  # -2
    x1_max = domain_max[1]  # 2
    x2_min = domain_min[2]  # -2
    x2_max = domain_max[2]  # 2

    for data in datas:
        (x0_l, x0_r, x1_l, x1_r, x2_l, x2_r,
         x01_w0, x01_w1, x01_b_l, x01_b,
         x02_w0, x02_w1, x02_b_l, x02_b,
         x12_w0, x12_w1, x12_b_l, x12_b,
         x0_k, x0_b_l, x0_b,
         x1_k, x1_b_l, x1_b,
         x2_k, x2_b_l, x2_b) = data

        x0_next = time_step * ((x0_k * x[0] + x0_b_l + var_x0_b)
                               + (x01_w0 * x[0] + x01_w1 * x[1] + x01_b_l + var_x01_b)
                               - (x02_w0 * x[0] + x02_w1 * x[2] + x02_b_l + var_x02_b)) + x[0]
        x1_next = time_step * (2 * (x01_w0 * x[0] + x01_w1 * x[1] + x01_b_l + var_x01_b)
                               + (x1_k * x[1] + x1_b_l + var_x1_b)) + x[1]
        x2_next = time_step * ((x12_w0 * x[1] + x12_w1 * x[2] + x12_b_l + var_x12_b)
                               - 2 * (x2_k * x[2] + x2_b_l + var_x2_b)) + x[2]

        f1_next, f2_next = express.getExpr([x0_next, x1_next, x2_next], model)

        # print("curvature of f1_next:", f1_next.curvature)
        # print("curvature of f2_next:", f2_next.curvature)

        u = cp.Variable()
        objective = cp.Maximize((f1_next - v) - (1 - alpha) * (u - f2))
        constraints = [v == f2_next, u == f1]
        constraints += [x0_l <= x[0], x[0] <= x0_r]
        constraints += [x1_l <= x[1], x[1] <= x1_r]
        constraints += [x2_l <= x[2], x[2] <= x2_r]
        constraints += [x0_min <= x0_next, x0_next <= x0_max]
        constraints += [x1_min <= x1_next, x1_next <= x1_max]
        constraints += [x2_min <= x2_next, x2_next <= x2_max]

        constraints += [0 <= var_x0_b, var_x0_b <= x0_b]
        constraints += [0 <= var_x1_b, var_x1_b <= x1_b]
        constraints += [0 <= var_x2_b, var_x2_b <= x2_b]
        constraints += [0 <= var_x01_b, var_x01_b <= x01_b]
        constraints += [0 <= var_x02_b, var_x02_b <= x02_b]
        constraints += [0 <= var_x12_b, var_x12_b <= x12_b]

        problem = cp.Problem(objective, constraints)
        # print("problem is DCCP:", dccp.is_dccp(problem))
        # x.value = np.array([1, 2])
        result = problem.solve(method='dccp')
        # print("x.value:", x.value)
        # print("var_b.value:", var_b.value)
        # print("problem.status:", problem.status)
        if problem.status == "Converged":
            # print("cost value =", result[0])
            # fw.write("================[" + str(x_l) + "," + str(x_r) + "]===================")
            # fw.write('\n')
            # fw.write(str(result[0]))
            # fw.write('\n')
            # fw.write('\n')

            if result[0] > 0.0:
                cex_domain.append(x.value)
                res_domain.append(result[0])
                print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] has cex point")
                fw.write(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] has cex point\n")
            else:
                print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is safe")
                fw.write(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is safe\n")
        else:
            print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is not converged")
            fw.write(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is not converged\n")

    fw.close()
    # print("==============cex and res=============")
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('../model/cdinn_c14_train.pt'), strict=True)
    # time_start_verify = time.time()
    verif(model)
    # time_end_verify = time.time()
    # print("verifing totally costs:", time_end_verify - time_start_verify)



